import json

from langchain.prompts import PromptTemplate
from langchain_ollama import OllamaLLM


# 1. 定义可用的函数
def get_weather(location: str) -> str:
    """获取指定位置的天气信息"""
    return f"模拟天气API: {location}的天气晴朗，温度25度"


def get_time(timezone: str = "UTC+8") -> str:
    """获取指定时区的时间"""
    return f"模拟时间API: {timezone}当前时间是下午3点"


# 2. 构建函数字典
available_functions = {
    "get_weather": get_weather,
    "get_time": get_time
}

# 3. 创建函数定义的描述
function_descriptions = [
    {
        "name": "get_weather",
        "description": "获取指定位置的天气信息",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "位置名称，如：北京、上海等"
                }
            },
            "required": ["location"]
        }
    },
    {
        "name": "get_time",
        "description": "获取指定时区的时间信息",
        "parameters": {
            "type": "object",
            "properties": {
                "timezone": {
                    "type": "string",
                    "description": "时区，如：UTC+8"
                }
            }
        }
    }
]

# 4. 创建提示模板
PROMPT_TEMPLATE = """你是一个助手，可以帮助用户获取信息。
你可以使用以下函数：

{function_descriptions}

用户的问题是：{user_input}

如果需要调用函数，请使用以下格式返回：
{{"function_call": {{"name": "函数名", "arguments": {{"参数名": "参数值"}}}}}}

如果不需要调用函数，直接返回回答即可。
"""

# 5. 初始化LLM
llm = OllamaLLM(
    model="qwen3:8b",
    temperature=0.7
)


def process_user_query(user_input: str) -> str:
    """处理用户输入并返回响应"""
    # 构建提示
    prompt = PromptTemplate(
        template=PROMPT_TEMPLATE,
        input_variables=["function_descriptions", "user_input"]
    )

    # 格式化函数描述
    formatted_descriptions = json.dumps(function_descriptions, ensure_ascii=False, indent=2)

    # 获取LLM响应
    response = llm.invoke(
        prompt.format(
            function_descriptions=formatted_descriptions,
            user_input=user_input
        )
    )

    try:
        # 尝试解析响应为JSON
        parsed_response = json.loads(response)

        # 检查是否需要调用函数
        if "function_call" in parsed_response:
            func_name = parsed_response["function_call"]["name"]
            func_args = parsed_response["function_call"]["arguments"]

            # 调用函数
            if func_name in available_functions:
                function_response = available_functions[func_name](**func_args)

                # 将函数调用结果再次发送给LLM进行处理
                second_prompt = f"""基于函数调用的结果：{function_response}
                请生成一个友好的回答给用户。
                原始问题是：{user_input}"""

                final_response = llm.invoke(second_prompt)
                return final_response
            else:
                return f"错误：未找到函数 {func_name}"
        else:
            return response
    except json.JSONDecodeError:
        # 如果响应不是JSON格式，直接返回原始响应
        return response


# 示例使用
if __name__ == "__main__":
    # 测试案例
    test_queries = [
        "北京现在天气怎么样？",
        "现在几点了？",
        "你是谁？"  # 这个问题不需要调用函数
    ]

    for query in test_queries:
        print(f"\n用户问题: {query}")
        print(f"助手回答: {process_user_query(query)}")
